/*
 * Copyright (C) 2008 by The Regents of the University of California
 * Redistribution of this file is permitted under the terms of the GNU
 * Public License (GPL).
 *
 * @author Junghoo "John" Cho <cho AT cs.ucla.edu>
 * @date 3/24/2008
 */
 
#include "BTreeIndex.h"
#include "BTreeNode.h"

using namespace std;

/*
 * BTreeIndex constructor
 */
BTreeIndex::BTreeIndex()
{
    rootPid = (int*)buffer;
    treeHeight = (((int*)buffer) + 1);
    nextFreePid = (((int*)buffer) + 2);
    branchingFactor = (((int*)buffer) + 3);
}

/*
 * Open the index file in read or write mode.
 * Under 'w' mode, the index file should be created if it does not exist.
 * @param indexname[IN] the name of the index file
 * @param mode[IN] 'r' for read, 'w' for write
 * @return error code. 0 if no error
 */
RC BTreeIndex::open(const string& indexname, char mode)
{
    RC rc = 0, canRead;
    if (rc = pf.open(indexname, mode) != 0)
    {
        fprintf(stderr, "Error: Could not open index file %s\n", indexname);
        return rc;
    }

    // Check if this file is new 
    if ((canRead = pf.read(0, buffer)) != 0)
    {
        // Initialize values if new
        *rootPid = 0;
        *treeHeight = 0;
        *nextFreePid = 1;
    }

    return rc;
}

/*
 * Close the index file.
 * @return error code. 0 if no error
 */
RC BTreeIndex::close()
{
    RC rc;
    rc = pf.close();
    return rc;
}

/*
 * Insert (key, RecordId) pair to the index.
 * @param key[IN] the key for the value inserted into the index
 * @param rid[IN] the RecordId for the record being inserted into the index
 * @return error code. 0 if no error
 */
RC BTreeIndex::insert(int key, const RecordId& rid)
{
    RC rc = 0;
    IndexCursor* cursor = new IndexCursor();

    // Initialize depth and non leaf node
    int currentDepth = 1;
    BTNonLeafNode* currentNonLeafNode = new BTNonLeafNode();

    // Go through non leafs until at leaf level
    PageId* pids = new PageId[*treeHeight];     // pid of leaf node and its ancestor nodes
    pids[*treeHeight-1] = *rootPid;
    int currentDepth = 1;
    for (; currentDepth < *treeHeight; ++currentDepth)
    {
        // Read node
        int currentIndex = *treeHeight - currentDepth;
        if ((rc = currentNonLeafNode->read(pids[currentIndex], pf)) != 0)
        {
            fprintf(stderr, "Error: Could not read non leaf node in page %d from page file.\n", pids[currentIndex]);
            delete currentNonLeafNode;
            return rc;
        }

        // Get next child node
        pids[currentIndex-1] = currentNonLeafNode->locateChildPtr(key, pids[currentIndex]);
    }

    // Delete non leaf node
    delete currentNonLeafNode;

    // Initialize leaf node
    BTLeafNode* leaf = new BTLeafNode();
    if ((rc = leaf->read(pids[0], pf)) != 0)
    {
        fprintf(stderr, "Error: Could not read leaf node in page %d from page file.\n", pids[0]);
        delete leaf;
        return rc;
    }
    int numKeys = leaf->getKeyCount();

    // Check if a new leaf needs to be allocated
    if (numKeys >= (*branchingFactor - 1))
    {
        int siblingKey = -1;
        BTLeafNode* sibling = new BTLeafNode();
        PageId siblingPid = *nextFreePid;
        
        // Increment next free page id
        *nextFreePid = *nextFreePid + 1;

        // Insert new key and split between old and new leaf nodes
        if ((rc = leaf->insertAndSplit(key, rid, *sibling, siblingKey)) != 0)
        {
            fprintf(stderr, "Error: Could not insert key %d and split between leaf nodes %d and %d.\n", key, pids[0], siblingKey);
            delete leaf;
            delete sibling;
            *nextFreePid = *nextFreePid - 1;
            return rc;
        }

        // Set new nextPageIds
        sibling->setNextNodePtr(leaf->getNextNodePtr());
        leaf->setNextNodePtr(*nextFreePid);

        // Add new pageId, key to parent node
        if ((rc = addPointerToParent(pids, currentDepth, siblingKey, siblingPid)) != 0)
        {
            fprintf(stderr, "Error: Could not insert key %d int parent node %d.\n", siblingKey, pids[1]);
            delete leaf;
            delete sibling;
            *nextFreePid = *nextFreePid - 1;
            return rc;
        }

        // Write changes to page file
        leaf->write(pids[0], pf);
        sibling->write(siblingPid, pf);

        // Delete sibling
        delete sibling;

        // Write table metadata to pagefile
        pf.write(0, buffer);
    }
    else
    {
        rc = leaf->insert(key, rid);
    }

    return rc;
}

/*
 * Find the leaf-node index entry whose key value is larger than or 
 * equal to searchKey, and output the location of the entry in IndexCursor.
 * IndexCursor is a "pointer" to a B+tree leaf-node entry consisting of
 * the PageId of the node and the SlotID of the index entry.
 * Note that, for range queries, we need to scan the B+tree leaf nodes.
 * For example, if the query is "key > 1000", we should scan the leaf
 * nodes starting with the key value 1000. For this reason,
 * it is better to return the location of the leaf node entry 
 * for a given searchKey, instead of returning the RecordId
 * associated with the searchKey directly.
 * Once the location of the index entry is identified and returned 
 * from this function, you should call readForward() to retrieve the
 * actual (key, rid) pair from the index.
 * @param key[IN] the key to find.
 * @param cursor[OUT] the cursor pointing to the first index entry
 *                    with the key value.
 * @return error code. 0 if no error.
 */
RC BTreeIndex::locate(int searchKey, IndexCursor& cursor)
{
    RC rc = 0;

    // Initialize depth and non leaf node
    int currentDepth = 1;
    BTNonLeafNode* currentNonLeafNode = new BTNonLeafNode();

    // Go through non leafs until at leaf level
    PageId childPid = *rootPid;
    for (int currentDepth = 1; currentDepth < *treeHeight; ++currentDepth)
    {
        // Read node
        if ((rc = currentNonLeafNode->read(childPid, pf)) != 0)
        {
            fprintf(stderr, "Error: Could not read non leaf node in page %d from page file.\n", childPid);
            delete currentNonLeafNode;
            return rc;
        }

        // Get next child node
        childPid = currentNonLeafNode->locateChildPtr(searchKey, childPid);
    }

    // Delete non leaf node
    delete currentNonLeafNode;

    // Initialize leaf node
    BTLeafNode* leaf = new BTLeafNode();
    if ((rc = leaf->read(childPid, pf)) != 0)
    {
        fprintf(stderr, "Error: Could not read leaf node in page %d from page file.\n", childPid);
        delete leaf;
        return rc;
    }

    // Set return values
    cursor.pid = childPid;
    rc = leaf->locate(searchKey, cursor.eid);

    // Delete leaf node
    delete leaf;

    return rc;
}

/*
 * Read the (key, rid) pair at the location specified by the index cursor,
 * and move foward the cursor to the next entry.
 * @param cursor[IN/OUT] the cursor pointing to an leaf-node index entry in the b+tree
 * @param key[OUT] the key stored at the index cursor location.
 * @param rid[OUT] the RecordId stored at the index cursor location.
 * @return error code. 0 if no error
 */
RC BTreeIndex::readForward(IndexCursor& cursor, int& key, RecordId& rid)
{
    RC rc = 0;
    int numEntries;

    // Ensure the cursor is using a valid PageId
    if (cursor.pid <= 0)
    {
        fprintf(stderr, "Error: Cursor pointing to invalid PageId %d. Probably end of tree.\n", cursor.pid);
        rc = RC_END_OF_TREE;
        return rc;
    }

    // Initialize the leaf
    BTLeafNode* leaf = new BTLeafNode();
    rc = leaf->read(cursor.pid, pf);

    // Ensure this entry exists
    if ((numEntries = leaf->getKeyCount()) <= cursor.eid)
    {
        fprintf(stderr, "Error: Cursor pointing to nonexistent entry %d in page %d. Node only has %d entries\n", cursor.eid, cursor.pid, numEntries);
        rc = RC_INVALID_CURSOR;
        return rc;
    }

    // Read key and rid into return values
    rc = leaf->readEntry(cursor.eid, key, rid);

    // Increment cursor to next entry
    if (cursor.eid < (numEntries - 1))
    {
        // There are still more entries in this node
        cursor.eid = cursor.eid + 1;
    }
    else
    {
        // Move cursor to next node
        cursor.eid = 0;
        cursor.pid = leaf->getNextNodePtr();
    }
    
    return rc;
}

/*
 * Recursively add the key, pid pair to a parent node, creating a new root
 * if necessary. Does not write out table metadata.
 * @param pids[IN] the PageId of the child node and its ancestor nodes
 * @param currentDepth[IN] the depth of the child node. (root depth = 1)
 * @param key[IN] the key to be added to the parent node.
 * @param pid[IN] the PageId to be added to the parent node.
 * @return error code. 0 if no error
 */
RC BTreeIndex::addPointerToParent(PageId pids[], int currentDepth, int key, PageId pid)
{
    // Check if a new root is needed
    if (currentDepth == 1)
    {
        PageId rootPid = *nextFreePid;

        // Increment next free page id
        *nextFreePid = *nextFreePid + 1;

        // Create new root
        BTNonLeafNode* root = new BTNonLeafNode();
        root->initializeRoot(pids[0], key, pid);

        // write new root out to page file
        root->write(*nextFreePid, pf);

        // Delete new root
        delete root;
    }
    else
    {
        // Need to add to parent node
        BTNonLeafNode* parent = new BTNonLeafNode();
        parent->read(pids[1], pf);

        // Check if a new non leaf node needs to be added
        if (parent->getKeyCount() >= *branchingFactor)
        {
            int midkey = -1;
            PageId nonLeafSiblingPid = *nextFreePid;

            // Increment next free page id
            *nextFreePid = *nextFreePid + 1;

            // Create new non leaf node
            BTNonLeafNode* nonLeafSibling = new BTNonLeafNode();
                
            // Insert new key/pid pair and split pairs between first node and sibling
            parent->insertAndSplit(key, pid, *nonLeafSibling, midkey);

            // Add new non leaf node to parent
            addPointerToParent(pids + 1, currentDepth - 1, key, nonLeafSiblingPid);

            // Write new non leaf sibling to page file
            nonLeafSibling->write(nonLeafSiblingPid, pf);

            // Delete new sibling non leaf node
            delete nonLeafSibling;
        }
        else
        {
            // Just add to parent node
            parent->insert(key, pid);
        }

        // Write parent node to file
        parent->write(pids[1], pf);
    }

    return 0;
}
